import torch
import model_input

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# input_data = model_input.get_input_data().to(device)
input_data=torch.randn(1, 3, 224, 224).to(device)

original_modal_path = '/home/recsche/tmp/new_layer_profile/results/formal_forward_graph.pth'

model_path_1 = '/home/recsche/tmp/new_layer_profile/results/modules/module_result0.pth'
model_path_2 = '/home/recsche/tmp/new_layer_profile/results/modules/module_result1.pth'
model_path_3 = '/home/recsche/tmp/new_layer_profile/results/modules/module_result2.pth'

original_model = torch.load(original_modal_path).to(device)

model_1 = torch.load(model_path_1).to(device)
model_2 = torch.load(model_path_2).to(device)
model_3 = torch.load(model_path_3).to(device)

original_model.eval()
model_1.eval()
model_2.eval()
model_3.eval()

output_1 = model_1(input_data)
output_2 = model_2(input_node1=output_1)
output_3 = model_3(input_node2=output_2)

output = original_model(input_data)

assert torch.allclose(output, output_3, atol=1e-6), "Outputs do not match!"
print("Outputs match! The graph was split correctly.")